"""
AI Historical Simulation Platform - Session Manager

This module provides comprehensive session management supporting 1000+ concurrent
users with efficient resource allocation, load balancing, and session persistence.
"""

import asyncio
import json
import logging
import time
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Set, Tuple
from dataclasses import dataclass, asdict
from collections import defaultdict
import weakref
from contextlib import asynccontextmanager

logger = logging.getLogger(__name__)


@dataclass
class UserSession:
    """Represents a user session with conversation state and metadata."""
    session_id: str
    user_id: Optional[str]
    created_at: datetime
    last_active: datetime
    active_figure: Optional[str]
    conversation_history: Dict[str, List[Dict]]  # figure_name -> messages
    session_metadata: Dict[str, Any]
    resource_usage: Dict[str, float]
    priority_level: int = 1  # 1=low, 2=medium, 3=high
    max_inactive_minutes: int = 30
    
    @property
    def is_expired(self) -> bool:
        """Check if session has expired due to inactivity."""
        inactive_time = datetime.now() - self.last_active
        return inactive_time.total_seconds() > (self.max_inactive_minutes * 60)
    
    @property
    def total_messages(self) -> int:
        """Get total number of messages in this session."""
        return sum(len(messages) for messages in self.conversation_history.values())
    
    def get_conversation_history(self, figure_name: str) -> List[Dict]:
        """Get conversation history for specific figure."""
        return self.conversation_history.get(figure_name, [])
    
    def add_message(self, figure_name: str, message: Dict[str, Any]) -> None:
        """Add message to conversation history."""
        if figure_name not in self.conversation_history:
            self.conversation_history[figure_name] = []
        
        self.conversation_history[figure_name].append({
            **message,
            'timestamp': datetime.now().isoformat()
        })
        
        self.last_active = datetime.now()
    
    def update_resource_usage(self, resource: str, value: float) -> None:
        """Update resource usage metrics."""
        self.resource_usage[resource] = value
        self.last_active = datetime.now()
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary representation."""
        return {
            'session_id': self.session_id,
            'user_id': self.user_id,
            'created_at': self.created_at.isoformat(),
            'last_active': self.last_active.isoformat(),
            'active_figure': self.active_figure,
            'total_messages': self.total_messages,
            'session_metadata': self.session_metadata,
            'resource_usage': self.resource_usage,
            'priority_level': self.priority_level,
            'is_expired': self.is_expired
        }


@dataclass
class SessionMetrics:
    """Session management performance metrics."""
    total_sessions: int
    active_sessions: int
    expired_sessions: int
    avg_session_duration: float
    total_messages: int
    messages_per_second: float
    resource_utilization: Dict[str, float]
    load_balance_efficiency: float
    concurrent_capacity_used: float


class SessionLoadBalancer:
    """Load balancer for distributing sessions across resources."""
    
    def __init__(self, max_sessions_per_worker: int = 100):
        self.max_sessions_per_worker = max_sessions_per_worker
        self.workers: Dict[str, Set[str]] = {}  # worker_id -> session_ids
        self.session_to_worker: Dict[str, str] = {}
        self.worker_metrics: Dict[str, Dict] = defaultdict(dict)
    
    def assign_session(self, session_id: str) -> str:
        """Assign session to least loaded worker."""
        # Find worker with minimum load
        min_load_worker = None
        min_load = float('inf')
        
        for worker_id, sessions in self.workers.items():
            if len(sessions) < min_load:
                min_load = len(sessions)
                min_load_worker = worker_id
        
        # Create new worker if needed
        if min_load_worker is None or min_load >= self.max_sessions_per_worker:
            worker_id = f"worker_{len(self.workers) + 1}"
            self.workers[worker_id] = set()
            min_load_worker = worker_id
        
        # Assign session
        self.workers[min_load_worker].add(session_id)
        self.session_to_worker[session_id] = min_load_worker
        
        return min_load_worker
    
    def remove_session(self, session_id: str) -> None:
        """Remove session from load balancer."""
        if session_id in self.session_to_worker:
            worker_id = self.session_to_worker[session_id]
            if worker_id in self.workers:
                self.workers[worker_id].discard(session_id)
            del self.session_to_worker[session_id]
    
    def get_worker_for_session(self, session_id: str) -> Optional[str]:
        """Get worker assignment for session."""
        return self.session_to_worker.get(session_id)
    
    def get_load_distribution(self) -> Dict[str, int]:
        """Get current load distribution across workers."""
        return {worker_id: len(sessions) for worker_id, sessions in self.workers.items()}
    
    def get_efficiency_score(self) -> float:
        """Calculate load balancing efficiency (0-1)."""
        if not self.workers:
            return 1.0
        
        loads = [len(sessions) for sessions in self.workers.values()]
        if not loads:
            return 1.0
        
        # Calculate coefficient of variation (lower is better)
        mean_load = sum(loads) / len(loads)
        if mean_load == 0:
            return 1.0
        
        variance = sum((load - mean_load) ** 2 for load in loads) / len(loads)
        cv = (variance ** 0.5) / mean_load
        
        # Convert to efficiency score (1 = perfectly balanced)
        return max(0, 1 - cv)


class SessionManager:
    """
    Advanced session manager supporting 1000+ concurrent users.
    
    Features:
    - Efficient session storage and retrieval
    - Automatic session cleanup and expiration
    - Load balancing across multiple workers
    - Resource monitoring and allocation
    - Session persistence and recovery
    - Priority-based resource allocation
    """
    
    def __init__(self, 
                 platform,
                 max_concurrent_sessions: int = 1000,
                 session_timeout_minutes: int = 30,
                 cleanup_interval_seconds: int = 60):
        """
        Initialize session manager.
        
        Args:
            platform: Reference to main platform
            max_concurrent_sessions: Maximum concurrent sessions
            session_timeout_minutes: Session timeout in minutes
            cleanup_interval_seconds: Cleanup interval in seconds
        """
        self.platform = platform  # Weak reference to avoid circular imports
        self.max_concurrent_sessions = max_concurrent_sessions
        self.session_timeout_minutes = session_timeout_minutes
        self.cleanup_interval_seconds = cleanup_interval_seconds
        
        # Session storage
        self.active_sessions: Dict[str, UserSession] = {}
        self.session_by_user: Dict[str, str] = {}  # user_id -> session_id
        
        # Load balancing
        self.load_balancer = SessionLoadBalancer(max_sessions_per_worker=100)
        
        # Resource tracking
        self.resource_pools: Dict[str, asyncio.Semaphore] = {
            'cpu_intensive': asyncio.Semaphore(50),
            'memory_intensive': asyncio.Semaphore(30),
            'io_intensive': asyncio.Semaphore(200)
        }
        
        # Performance metrics
        self.metrics = {
            'sessions_created': 0,
            'sessions_expired': 0,
            'total_messages': 0,
            'peak_concurrent_sessions': 0,
            'avg_session_duration': 0.0,
            'resource_waits': defaultdict(int)
        }
        
        # Background tasks
        self.cleanup_task: Optional[asyncio.Task] = None
        self.metrics_task: Optional[asyncio.Task] = None
        self.shutdown_event = asyncio.Event()
        
        logger.info(f"SessionManager initialized for {max_concurrent_sessions} concurrent sessions")
    
    async def initialize(self) -> None:
        """Initialize session manager and start background tasks."""
        logger.info("Initializing SessionManager")
        
        # Start background tasks
        self.cleanup_task = asyncio.create_task(self._background_cleanup())
        self.metrics_task = asyncio.create_task(self._background_metrics())
        
        logger.info("SessionManager initialized successfully")
    
    async def shutdown(self) -> None:
        """Shutdown session manager gracefully."""
        logger.info("Shutting down SessionManager")
        
        # Signal shutdown
        self.shutdown_event.set()
        
        # Cancel background tasks
        if self.cleanup_task:
            self.cleanup_task.cancel()
            try:
                await self.cleanup_task
            except asyncio.CancelledError:
                pass
        
        if self.metrics_task:
            self.metrics_task.cancel()
            try:
                await self.metrics_task
            except asyncio.CancelledError:
                pass
        
        # Save session state
        await self._save_session_state()
        
        # Clear resources
        self.active_sessions.clear()
        self.session_by_user.clear()
        
        logger.info("SessionManager shutdown completed")
    
    async def create_session(self, 
                           user_id: Optional[str] = None,
                           session_metadata: Optional[Dict[str, Any]] = None,
                           priority_level: int = 1) -> str:
        """
        Create new user session.
        
        Args:
            user_id: Optional user identifier
            session_metadata: Optional session metadata
            priority_level: Session priority (1=low, 2=medium, 3=high)
            
        Returns:
            New session ID
            
        Raises:
            RuntimeError: If concurrent session limit exceeded
        """
        # Check concurrent session limit
        if len(self.active_sessions) >= self.max_concurrent_sessions:
            # Try to clean up expired sessions first
            await self._cleanup_expired_sessions()
            
            if len(self.active_sessions) >= self.max_concurrent_sessions:
                raise RuntimeError(f"Maximum concurrent sessions ({self.max_concurrent_sessions}) exceeded")
        
        # Generate unique session ID
        session_id = f"session_{int(time.time() * 1000)}_{len(self.active_sessions)}"
        
        # Create session
        session = UserSession(
            session_id=session_id,
            user_id=user_id,
            created_at=datetime.now(),
            last_active=datetime.now(),
            active_figure=None,
            conversation_history={},
            session_metadata=session_metadata or {},
            resource_usage={},
            priority_level=priority_level
        )
        
        # Store session
        self.active_sessions[session_id] = session
        
        if user_id:
            self.session_by_user[user_id] = session_id
        
        # Assign to load balancer
        worker_id = self.load_balancer.assign_session(session_id)
        session.session_metadata['worker_id'] = worker_id
        
        # Update metrics
        self.metrics['sessions_created'] += 1
        self.metrics['peak_concurrent_sessions'] = max(
            self.metrics['peak_concurrent_sessions'],
            len(self.active_sessions)
        )
        
        logger.debug(f"Created session {session_id} for user {user_id} on worker {worker_id}")
        return session_id
    
    async def get_session(self, session_id: str) -> Optional[UserSession]:
        """
        Get session by ID.
        
        Args:
            session_id: Session identifier
            
        Returns:
            User session or None if not found/expired
        """
        session = self.active_sessions.get(session_id)
        
        if session and session.is_expired:
            await self._remove_session(session_id, reason="expired")
            return None
        
        return session
    
    async def get_session_by_user(self, user_id: str) -> Optional[UserSession]:
        """
        Get session by user ID.
        
        Args:
            user_id: User identifier
            
        Returns:
            User session or None if not found
        """
        session_id = self.session_by_user.get(user_id)
        if session_id:
            return await self.get_session(session_id)
        return None
    
    async def update_session_activity(self, 
                                    session_id: str,
                                    figure_name: str,
                                    user_input: str) -> None:
        """
        Update session activity with new interaction.
        
        Args:
            session_id: Session identifier
            figure_name: Name of historical figure
            user_input: User's input message
        """
        session = await self.get_session(session_id)
        if not session:
            raise ValueError(f"Session {session_id} not found")
        
        # Add message to conversation history
        message = {
            'type': 'user_input',
            'content': user_input,
            'figure': figure_name
        }
        
        session.add_message(figure_name, message)
        session.active_figure = figure_name
        
        # Update metrics
        self.metrics['total_messages'] += 1
    
    async def add_figure_response(self,
                                session_id: str,
                                figure_name: str,
                                response: str,
                                response_metadata: Optional[Dict[str, Any]] = None) -> None:
        """
        Add figure response to session.
        
        Args:
            session_id: Session identifier
            figure_name: Name of historical figure
            response: Figure's response
            response_metadata: Optional response metadata
        """
        session = await self.get_session(session_id)
        if not session:
            raise ValueError(f"Session {session_id} not found")
        
        # Add response to conversation history
        message = {
            'type': 'figure_response',
            'content': response,
            'figure': figure_name,
            'metadata': response_metadata or {}
        }
        
        session.add_message(figure_name, message)
        
        # Update metrics
        self.metrics['total_messages'] += 1
    
    async def remove_session(self, session_id: str) -> bool:
        """
        Remove session manually.
        
        Args:
            session_id: Session identifier
            
        Returns:
            True if session was removed
        """
        return await self._remove_session(session_id, reason="manual")
    
    async def get_active_sessions(self) -> List[str]:
        """Get list of active session IDs."""
        # Clean up expired sessions first
        await self._cleanup_expired_sessions()
        return list(self.active_sessions.keys())
    
    async def get_session_count_by_figure(self) -> Dict[str, int]:
        """Get count of sessions by active figure."""
        figure_counts = defaultdict(int)
        
        for session in self.active_sessions.values():
            if session.active_figure and not session.is_expired:
                figure_counts[session.active_figure] += 1
        
        return dict(figure_counts)
    
    async def get_metrics(self) -> SessionMetrics:
        """Get current session metrics."""
        active_count = len([s for s in self.active_sessions.values() if not s.is_expired])
        expired_count = len([s for s in self.active_sessions.values() if s.is_expired])
        
        # Calculate average session duration
        total_duration = 0
        session_count = 0
        
        for session in self.active_sessions.values():
            duration = (session.last_active - session.created_at).total_seconds()
            total_duration += duration
            session_count += 1
        
        avg_duration = total_duration / session_count if session_count > 0 else 0
        
        # Calculate messages per second
        uptime = (datetime.now() - getattr(self, 'start_time', datetime.now())).total_seconds()
        messages_per_second = self.metrics['total_messages'] / uptime if uptime > 0 else 0
        
        # Calculate resource utilization
        resource_utilization = {}
        for resource_name, semaphore in self.resource_pools.items():
            available = semaphore._value
            total_capacity = semaphore._initial_value
            utilization = (total_capacity - available) / total_capacity
            resource_utilization[resource_name] = utilization
        
        return SessionMetrics(
            total_sessions=len(self.active_sessions),
            active_sessions=active_count,
            expired_sessions=expired_count,
            avg_session_duration=avg_duration,
            total_messages=self.metrics['total_messages'],
            messages_per_second=messages_per_second,
            resource_utilization=resource_utilization,
            load_balance_efficiency=self.load_balancer.get_efficiency_score(),
            concurrent_capacity_used=len(self.active_sessions) / self.max_concurrent_sessions
        )
    
    @asynccontextmanager
    async def acquire_resource(self, resource_type: str, session_id: str):
        """
        Context manager for acquiring resources with session tracking.
        
        Args:
            resource_type: Type of resource to acquire
            session_id: Session requesting resource
        """
        if resource_type not in self.resource_pools:
            raise ValueError(f"Unknown resource type: {resource_type}")
        
        semaphore = self.resource_pools[resource_type]
        
        # Track wait time
        wait_start = time.time()
        async with semaphore:
            wait_time = time.time() - wait_start
            
            if wait_time > 0.1:  # Log if wait > 100ms
                self.metrics['resource_waits'][resource_type] += 1
                logger.debug(f"Session {session_id} waited {wait_time:.3f}s for {resource_type}")
            
            # Update session resource usage
            session = self.active_sessions.get(session_id)
            if session:
                session.update_resource_usage(resource_type, time.time())
            
            try:
                yield
            finally:
                # Resource automatically released by semaphore
                pass
    
    async def get_load_distribution(self) -> Dict[str, Any]:
        """Get current load distribution information."""
        return {
            'worker_distribution': self.load_balancer.get_load_distribution(),
            'efficiency_score': self.load_balancer.get_efficiency_score(),
            'total_workers': len(self.load_balancer.workers),
            'sessions_per_worker': {
                worker_id: len(sessions) 
                for worker_id, sessions in self.load_balancer.workers.items()
            }
        }
    
    # Private methods
    
    async def _remove_session(self, session_id: str, reason: str = "unknown") -> bool:
        """Remove session from all tracking structures."""
        session = self.active_sessions.get(session_id)
        if not session:
            return False
        
        # Remove from active sessions
        del self.active_sessions[session_id]
        
        # Remove user mapping
        if session.user_id and session.user_id in self.session_by_user:
            del self.session_by_user[session.user_id]
        
        # Remove from load balancer
        self.load_balancer.remove_session(session_id)
        
        # Update metrics
        if reason == "expired":
            self.metrics['sessions_expired'] += 1
        
        logger.debug(f"Removed session {session_id} (reason: {reason})")
        return True
    
    async def _cleanup_expired_sessions(self) -> int:
        """Clean up expired sessions."""
        expired_sessions = [
            session_id for session_id, session in self.active_sessions.items()
            if session.is_expired
        ]
        
        for session_id in expired_sessions:
            await self._remove_session(session_id, reason="expired")
        
        if expired_sessions:
            logger.debug(f"Cleaned up {len(expired_sessions)} expired sessions")
        
        return len(expired_sessions)
    
    async def _background_cleanup(self) -> None:
        """Background task for periodic session cleanup."""
        while not self.shutdown_event.is_set():
            try:
                await self._cleanup_expired_sessions()
                await asyncio.sleep(self.cleanup_interval_seconds)
            except Exception as e:
                logger.error(f"Background cleanup error: {e}")
                await asyncio.sleep(self.cleanup_interval_seconds * 2)
    
    async def _background_metrics(self) -> None:
        """Background task for metrics collection."""
        self.start_time = datetime.now()
        
        while not self.shutdown_event.is_set():
            try:
                # Update metrics
                metrics = await self.get_metrics()
                
                # Log performance warnings
                if metrics.concurrent_capacity_used > 0.8:
                    logger.warning(f"High session capacity usage: {metrics.concurrent_capacity_used:.1%}")
                
                if metrics.load_balance_efficiency < 0.7:
                    logger.warning(f"Poor load balance efficiency: {metrics.load_balance_efficiency:.2f}")
                
                # Check resource utilization
                for resource, utilization in metrics.resource_utilization.items():
                    if utilization > 0.9:
                        logger.warning(f"High {resource} utilization: {utilization:.1%}")
                
                await asyncio.sleep(30)  # Check every 30 seconds
                
            except Exception as e:
                logger.error(f"Background metrics error: {e}")
                await asyncio.sleep(60)
    
    async def _save_session_state(self) -> None:
        """Save session state for recovery."""
        try:
            session_data = {
                'sessions': {
                    session_id: session.to_dict() 
                    for session_id, session in self.active_sessions.items()
                },
                'metrics': self.metrics,
                'timestamp': datetime.now().isoformat()
            }
            
            # Save to file (in production, would use database)
            with open('session_state.json', 'w') as f:
                json.dump(session_data, f, indent=2, default=str)
            
            logger.info(f"Saved state for {len(self.active_sessions)} sessions")
            
        except Exception as e:
            logger.error(f"Failed to save session state: {e}")
    
    async def load_session_state(self, filepath: str = 'session_state.json') -> None:
        """Load session state for recovery."""
        try:
            with open(filepath, 'r') as f:
                session_data = json.load(f)
            
            # Restore sessions (only recent ones)
            cutoff_time = datetime.now() - timedelta(hours=1)
            restored_count = 0
            
            for session_id, session_dict in session_data.get('sessions', {}).items():
                created_at = datetime.fromisoformat(session_dict['created_at'])
                
                if created_at > cutoff_time:
                    # Recreate session object
                    session = UserSession(
                        session_id=session_dict['session_id'],
                        user_id=session_dict['user_id'],
                        created_at=created_at,
                        last_active=datetime.fromisoformat(session_dict['last_active']),
                        active_figure=session_dict['active_figure'],
                        conversation_history=session_dict.get('conversation_history', {}),
                        session_metadata=session_dict.get('session_metadata', {}),
                        resource_usage=session_dict.get('resource_usage', {}),
                        priority_level=session_dict.get('priority_level', 1)
                    )
                    
                    self.active_sessions[session_id] = session
                    
                    if session.user_id:
                        self.session_by_user[session.user_id] = session_id
                    
                    # Re-assign to load balancer
                    self.load_balancer.assign_session(session_id)
                    
                    restored_count += 1
            
            # Restore metrics
            if 'metrics' in session_data:
                self.metrics.update(session_data['metrics'])
            
            logger.info(f"Restored {restored_count} sessions from saved state")
            
        except FileNotFoundError:
            logger.info("No session state file found, starting fresh")
        except Exception as e:
            logger.error(f"Failed to load session state: {e}")


# Utility functions for session management

async def get_or_create_session(session_manager: SessionManager,
                               session_id: Optional[str] = None,
                               user_id: Optional[str] = None,
                               **kwargs) -> Tuple[str, UserSession]:
    """
    Get existing session or create new one.
    
    Args:
        session_manager: Session manager instance
        session_id: Optional existing session ID
        user_id: Optional user ID
        **kwargs: Additional arguments for session creation
        
    Returns:
        Tuple of (session_id, session)
    """
    # Try to get existing session
    if session_id:
        session = await session_manager.get_session(session_id)
        if session:
            return session_id, session
    
    # Try to get session by user ID
    if user_id:
        session = await session_manager.get_session_by_user(user_id)
        if session:
            return session.session_id, session
    
    # Create new session
    new_session_id = await session_manager.create_session(user_id=user_id, **kwargs)
    session = await session_manager.get_session(new_session_id)
    
    return new_session_id, session


def calculate_session_priority(user_metadata: Dict[str, Any]) -> int:
    """
    Calculate session priority based on user metadata.
    
    Args:
        user_metadata: User metadata dictionary
        
    Returns:
        Priority level (1=low, 2=medium, 3=high)
    """
    # Example priority calculation logic
    if user_metadata.get('subscription_tier') == 'premium':
        return 3
    elif user_metadata.get('is_educator'):
        return 2
    else:
        return 1